#include "cartographer/rcm/imu/initialize.h"

namespace cartographer
{
    ImuInitialize::ImuInitialize(size_t win_min, double duration, double motion_var)
        : win_min_(win_min),
          duration_(duration),
          motion_var_(motion_var)
    {
    }

    void ImuInitialize::AddData(const sensor::ImuData &data)
    {
        if (inited_)
            return;
        if (skip_count_ < skip_times_)
        {
            skip_count_++;
            return;
        }
        const auto &stamp = common::ToUniversal(data.time);
        if (HasNan(data))
        {
            LOG(INFO) << "imu HasNan: " << stamp << "\n";
            return;
        }
        int deque_size = deque_.size();
        if (deque_size < 1)
        {
            deque_.push_back(data);
            return;
        }
        const auto &last_time = deque_.back().time;
        if (data.time <= last_time)
        {
            LOG(INFO) << "imu disorder: " << stamp << "\n";
            return;
        }
        deque_.push_back(data);
        if (int(deque_size) < 2 * int(win_min_) - 1)
            return;
        InitializeWindow();
    }

    void ImuInitialize::PopFront(size_t size)
    {
        size_t count = 0;
        while (count < size)
        {
            count++;
            deque_.pop_front();
        }
    }

    void ImuInitialize::CalcAverage(Eigen::Vector3d &gyro_avg, Eigen::Vector3d &accel_avg, const ImuSet &set)
    {
        gyro_avg = Eigen::Vector3d::Zero();
        accel_avg = Eigen::Vector3d::Zero();
        for (auto &temp : set)
        {
            gyro_avg += temp.angular_velocity;
            accel_avg += temp.linear_acceleration;
            //LOG(INFO) << "imu " << common::ToUniversal(temp.time) << "\n accel: " << VecToStr(temp.linear_acceleration) << " gyro: " << VecToStr(temp.angular_velocity) << "\n";
        }
        gyro_avg /= int(set.size());
        accel_avg /= int(set.size());
        LOG(INFO) << "avg accel: " << VecToStr(accel_avg) << " gyro: " << VecToStr(gyro_avg) << "\n";
    }

    bool ImuInitialize::HasMotion(const std::vector<sensor::ImuData> &set, const Eigen::Vector3d &accel_avg)
    {
        double acc_var = 0;
        for (auto &temp : set)
        {
            acc_var += (temp.linear_acceleration - accel_avg).dot(temp.linear_acceleration - accel_avg);
        }
        acc_var = std::sqrt(acc_var / (int)set.size());
        LOG(INFO) << "acc_var: " << acc_var << "\n";
        return acc_var > motion_var_;
    }

    void ImuInitialize::InitializeWindow()
    {
        auto &first_time = deque_.front().time;
        auto &last_time = deque_.back().time;
        auto elapsed = common::ToSeconds(last_time - first_time);
        if (elapsed < 2.0 * duration_)
        {
            //LOG(INFO) << "elapsed small " << elapsed << "\n";
            return;
        }
        std::vector<sensor::ImuData> first_win;
        std::vector<sensor::ImuData> second_win;
        for (auto iter = deque_.begin(); iter < deque_.end(); iter++)
        {
            //int mark = 0;
            auto elapsed = common::ToSeconds(iter->time - first_time);
            if (elapsed >= 0 && elapsed < duration_)
            {
                //mark = 1;
                first_win.push_back(*iter);
            }
            if (elapsed >= duration_ && elapsed < 2.0 * duration_)
            {
                //mark = 2;
                second_win.push_back(*iter);
            }
            //LOG(INFO) << "elapsed: " << elapsed << " mark " << mark << "\n";
        }
        auto first_size = first_win.size();
        auto second_size = second_win.size();
        LOG(INFO) << "windows: " << deque_.size() << " " << first_size << " " << second_size << "\n";
        if (first_size < win_min_ || second_size < win_min_)
        {
            if (first_size > 0)
                PopFront(first_size);
            return;
        }
        // 计算第二段加速度均值
        Eigen::Vector3d gyro_avg = Eigen::Vector3d::Zero();
        Eigen::Vector3d accel_avg = Eigen::Vector3d::Zero();
        LOG(INFO) << "second_win avg:\n";
        CalcAverage(gyro_avg, accel_avg, second_win);
        // 根据均值计算第二段加速度方差
        if (!HasMotion(second_win, accel_avg))
        {
            // 第二段处于静止态
            PopFront(first_size);
            return;
        }
        // 计算第一段加速度均值
        LOG(INFO) << "first_win avg:\n";
        CalcAverage(gyro_avg, accel_avg, first_win);
        // 根据均值计算第一段加速度方差
        if (HasMotion(first_win, accel_avg))
        {
            // 第一段有运动
            PopFront(first_size);
            PopFront(second_size);
            return;
        }
        CalcBias(gyro_avg, accel_avg);
    }

    void ImuInitialize::CalcBias(const Eigen::Vector3d &gyro_avg, const Eigen::Vector3d &accel_avg)
    {
        Eigen::Vector3d g_vec(0.0, 0.0, -9.80665);
        Eigen::Quaterniond rotate = Eigen::Quaternion<double>::FromTwoVectors(accel_avg, -g_vec);
        Eigen::Vector3d acc_bias = rotate.conjugate().toRotationMatrix() * g_vec + accel_avg;
        Eigen::Vector3d gyro_bias = gyro_avg;
        imu_bias_ = ImuBias(acc_bias, gyro_bias);
        inited_ = true;
        LOG(INFO) << "imu_bias acc: " << VecToStr(acc_bias) << " "
                  << "gyro:" << VecToStr(gyro_bias) << "\n";
    }

    void ImuInitialize::InitializeStatic()
    {
        int num_readings = 0;
        Eigen::Vector3d accel_accum;
        Eigen::Vector3d gyro_accum;
        accel_accum.setZero();
        gyro_accum.setZero();
        LOG(INFO) << "imu init " << deque_.size() << "\n";
        for (auto iter = deque_.begin(); iter < deque_.end(); iter++)
        {
            Print(*iter);
            num_readings++;
            gyro_accum += iter->angular_velocity;
            accel_accum += iter->linear_acceleration;
        }
        Eigen::Vector3d gyro_avg = gyro_accum / num_readings;
        Eigen::Vector3d accel_avg = accel_accum / num_readings;
        CalcBias(gyro_avg, accel_avg);
    }
}